import torch
from torch import nn
from torch.utils.data import TensorDataset, Dataset, DataLoader
from torch.optim import SGD, Adam
device = 'cuda' if torch.cuda.is_available() else 'cpu'
from torchvision import datasets
import numpy as np
import matplotlib.pyplot as plt

X_train = torch.tensor([
    [[[1,2,3,4],[2,3,4,5],[5,6,7,8],[1,3,4,5]]],
    [[[-1,2,3,-4],[2,-3,4,5],[-5,6,-7,8],[-1,-3,-4,-5]]]
]).to(device).float()
X_train /= 8
y_train = torch.tensor([0,1]).to(device).float()

model = torch.load('cnn01.pth')

# print(list(model.children()))
# for pp in model.children():
#     print(pp)
#
(conn_w,conn_b),(lin_w,lin_b) =[(layer.weight.data,layer.bias.data) for layer in list(model.children()) if hasattr(layer,'weight')]
# print(conn_w)
#
h_im,w_im = X_train.shape[2:]
print(X_train.shape[2:])
h_conv ,w_conv = conn_w.shape[2:]
# print(conn_w.shape[2:])
sumprod = torch.zeros((h_im - h_conv + 1,w_im,w_conv + 1))
#
for i in range(h_im -h_conv + 1):
    for j in range(w_im - w_conv + 1):
        img_subset = X_train[0,0,i:(i+1),j:(j+3)]
        model_filter = conn_w.reshape(3,3)
        val = torch.sum(img_subset * model_filter) + conn_b
        sumprod[i,j] = val

print(sumprod)